﻿using Assets.Scripts.Bvh.Common;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;
using UnityEngine;

namespace Assets.Scripts.Bvh.Gpu
{
    class GpuSorter
    {
        public ComputeShader bitonicSortCs;
        public ComputeBuffer leafNodesBuffer;

        public void BitonicSort(int rawCount)
        {
            var expandKernel = bitonicSortCs.FindKernel("ExpandData");
            var targetSortBuffer = leafNodesBuffer;
            var expanded = false;
            var dataCount = rawCount;

            //判断长度
            if (!Mathf.IsPowerOfTwo(dataCount))
            {
                var expandLength = Mathf.NextPowerOfTwo(dataCount);
                Debug.Log("Current data count : " + dataCount + "__expand length : " + expandLength);
                var expandBuffer = new ComputeBuffer(expandLength, Marshal.SizeOf<RadixTreeNode>());

                bitonicSortCs.SetInt("_DataLength", dataCount);
                bitonicSortCs.SetInt("_ExpandLength", expandLength);
                bitonicSortCs.SetBuffer(expandKernel, "_ExpandNodesBuffer", expandBuffer);
                bitonicSortCs.SetBuffer(expandKernel, "_LeafNodes", leafNodesBuffer);

                var dispatch = Mathf.CeilToInt(expandLength / 1024.0f);
                bitonicSortCs.Dispatch(expandKernel, dispatch, 1, 1);


                targetSortBuffer = expandBuffer;
                dataCount = expandLength;

                expanded = true;
            }

            var genSeqKernel = bitonicSortCs.FindKernel("GenerateSequence");
            var sortKernel = bitonicSortCs.FindKernel("BitonicSort");

            bitonicSortCs.SetBuffer(genSeqKernel, "_LeafNodes", targetSortBuffer);
            bitonicSortCs.SetInt("_DataLength", dataCount);
            var dispatchX = Mathf.CeilToInt(dataCount / 1024.0f);

            //双调序列生成
            {
                var signLength = 2;
                while (signLength < dataCount)
                {
                    var compareLength = signLength / 2;
                    bitonicSortCs.SetInt("_SignLength", signLength);

                    while (compareLength > 0)
                    {
                        bitonicSortCs.SetInt("_CompareLength", compareLength);
                        bitonicSortCs.Dispatch(genSeqKernel, dispatchX, 1, 1);
                        compareLength /= 2;
                    }
                    signLength *= 2;
                }
            }

            //双调排序
            {
                bitonicSortCs.SetBuffer(sortKernel, "_LeafNodes", targetSortBuffer);
                bitonicSortCs.SetInt("_DataLength", dataCount);

                var compareLength = dataCount / 2;
                while (compareLength != 0)
                {
                    bitonicSortCs.SetInt("_CompareLength", compareLength);
                    bitonicSortCs.Dispatch(sortKernel, dispatchX, 1, 1);

                    compareLength /= 2;
                }
            }

            //如果扩容了，再剔除一下。
            if (expanded)
            {
                var retrieveKernel = bitonicSortCs.FindKernel("RetrieveData");
                bitonicSortCs.SetInt("_DataLength", rawCount);
                bitonicSortCs.SetInt("_ExpandLength", dataCount);
                bitonicSortCs.SetBuffer(retrieveKernel, "_ExpandNodesBuffer", targetSortBuffer);
                bitonicSortCs.SetBuffer(retrieveKernel, "_LeafNodes", leafNodesBuffer);

                var dispatch = Mathf.CeilToInt(rawCount / 1024.0f);

                bitonicSortCs.Dispatch(retrieveKernel, dispatch, 1, 1);

                targetSortBuffer?.Dispose();
            }
        }
    }
}
